# 输入：文件路径
# 输出：vtkActor
import vtkmodules.all as vtk


def to_polydata(data_input, type='structured'):
    if type == 'structured':
        # 创建一个几何过滤器，将非结构化网格数据转换为多边形数据
        geometryFilter = vtk.vtkGeometryFilter()
        geometryFilter.SetInputData(data_input)
        geometryFilter.Update()
        # 获取转换后的多边形数据
        polyData = geometryFilter.GetOutput()
        return polyData
    else:
        ifilter = vtk.vtkDataSetSurfaceFilter()
        ifilter.SetInputData(data_input)
        ifilter.Update()
        polyData = ifilter.GetOutput()
        return polyData


class IVTKReader:
    def __init__(self, file_type=""):
        self.type = file_type
        if self.type == "unstructured":
            self.v_reader = vtk.vtkUnstructuredGridReader()
        if self.type == "structured":
            self.v_reader = vtk.vtkStructuredGridReader()
        if self.type == "dataset":
            self.v_reader = vtk.vtkDataSetReader()
            self.v_reader.ReadAllColorScalarsOn()
            self.v_reader.ReadAllVectorsOn()
        else:
            self.v_reader = vtk.vtkPolyDataReader()

    def read(self, file_path, force_update=True):
        self.v_reader.SetFileName(file_path)
        if force_update:
            self.v_reader.Update()

    def output(self):
        # 非结构化
        if self.type == "unstructured":
            unstructuredGrid = self.v_reader.GetOutput()
            return to_polydata(unstructuredGrid)
        # 结构化
        if self.type == "structured":
            structured_grid = self.v_reader.GetOutput()
            return to_polydata(structured_grid)
        # 数据集 注意这里返回的不是 polydata #️⃣
        if self.type == "dataset":
            data_set = self.v_reader.GetOutput()
            # 根据数据类型选择映射器
            if data_set.IsA("vtkStructuredGrid"):
                return to_polydata(data_set)
            elif data_set.IsA("vtkUnstructuredGrid"):
                return to_polydata(data_set, "vtkUnstructuredGrid")
            elif data_set.IsA("vtkPolyData"):
                return data_set
        else:
            # 多边形数据
            return self.v_reader.GetOutput()
